{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import Input, Dense, MultiHeadAttention, LayerNormalization, Dropout, Embedding, Flatten\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_next(x):\n",
    "    e = {'+':0,'-':0,'*':1,'/':1}\n",
    "    current_order = 0\n",
    "    operators = []\n",
    "    parameters = []\n",
    "    orders = []\n",
    "    s = ''\n",
    "    for i,j in enumerate(x):\n",
    "        if j=='(':\n",
    "            current_order += 2\n",
    "            s += j\n",
    "        elif j==')':\n",
    "            current_order -= 2\n",
    "            s += j\n",
    "        elif j in e:\n",
    "            operators.append(j)\n",
    "            orders.append(current_order+e[j])\n",
    "            parameters.append(s)\n",
    "            s = ''\n",
    "        else:\n",
    "            s += j\n",
    "    parameters.append(s)\n",
    "    if len(operators)<=1:\n",
    "        return x\n",
    "    else:\n",
    "        d = np.argmax(orders)\n",
    "        a = parameters[d].replace('(','').replace(')','')\n",
    "        b = parameters[d+1].replace('(','').replace(')','')\n",
    "        p = f'{a}{operators[d]}{b}'\n",
    "        c = x.replace(p,'&').replace('(&)','&')\n",
    "        return f'{p}|{c}'\n",
    "_y = \"\"\n",
    "def parse(x):\n",
    "    global _y\n",
    "    result = parse_next(x)\n",
    "    if '|' not in result:\n",
    "        result = str(eval(x))\n",
    "        #print(result)\n",
    "        _y += \" \"+x\n",
    "        return result\n",
    "    else:\n",
    "        #print(result)\n",
    "        _y += \" \"+result\n",
    "        a,b = result.split('|')\n",
    "        c = str(eval(a))\n",
    "        d = b.replace('&',c)\n",
    "        parse(d)\n",
    "tmp = []\n",
    "X = []\n",
    "Y = []\n",
    "def generate_expression(max_depth=3,max_value=10):\n",
    "    if max_depth == 0:\n",
    "        return str(np.random.randint(0,max_value))\n",
    "    else:\n",
    "        operators = ['+','*']\n",
    "        operator = np.random.choice(operators)\n",
    "        left_expr = generate_expression(max_depth-1, max_value)\n",
    "        right_expr = generate_expression(max_depth-1, max_value)\n",
    "        if operator in ['*','-'] or np.random.choice([True,False]):\n",
    "            result = f\"({left_expr}{operator}{right_expr})\"\n",
    "        else:\n",
    "            result = f\"{left_expr}{operator}{right_expr}\"\n",
    "        return result\n",
    "tokenizer = tf.keras.preprocessing.text.Tokenizer(char_level=True)\n",
    "tokenizer.fit_on_texts(['0123456789+*()|&'])\n",
    "def data_generator(batchsize):\n",
    "    global X,Y,_y\n",
    "    while 1:\n",
    "        if len(X)<batchsize:\n",
    "            try:\n",
    "                er = 3\n",
    "                X.append(generate_expression(er,100))\n",
    "                parse(X[-1])\n",
    "                Y.append(_y[1:])\n",
    "                _y = \"\"\n",
    "            except:\n",
    "                X = []\n",
    "                Y = []\n",
    "                _y = \"\"\n",
    "        else:\n",
    "            a = X[:batchsize]\n",
    "            b = Y[:batchsize]\n",
    "            X = X[batchsize:]\n",
    "            Y = Y[batchsize:]\n",
    "            a = tf.keras.preprocessing.sequence.pad_sequences(tokenizer.texts_to_sequences(a),padding='post',maxlen=100)\n",
    "            b = tf.keras.preprocessing.sequence.pad_sequences(tokenizer.texts_to_sequences(b),padding='post',maxlen=300)\n",
    "            X = []\n",
    "            Y = []\n",
    "            yield [a,b]\n",
    "data_g = data_generator(32)\n",
    "def compute(equation,debug=False):\n",
    "    result = model.predict(tf.keras.preprocessing.sequence.pad_sequences(tokenizer.texts_to_sequences([equation]),padding='post',maxlen=100),verbose=0)\n",
    "    result = tokenizer.sequences_to_texts(np.argmax(result,axis=-1))[0].split(' ')\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(tf.keras.layers.Layer):\n",
    "    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)\n",
    "        self.ffn = tf.keras.Sequential([Dense(ff_dim, activation=\"relu\"), Dense(embed_dim),])\n",
    "        self.layernorm1 = LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = LayerNormalization(epsilon=1e-6)\n",
    "        self.dropout1 = Dropout(rate)\n",
    "        self.dropout2 = Dropout(rate)\n",
    "    def call(self,inputs,training):\n",
    "        attn_output = self.att(inputs,inputs)\n",
    "        attn_output = self.dropout1(attn_output,training=training)\n",
    "        out1 = self.layernorm1(inputs+attn_output)\n",
    "        ffn_output = self.ffn(out1)\n",
    "        ffn_output = self.dropout2(ffn_output,training=training)\n",
    "        return self.layernorm2(out1+ffn_output)\n",
    "class TransformerBlock2(tf.keras.layers.Layer):\n",
    "    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):\n",
    "        super(TransformerBlock2, self).__init__()\n",
    "        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)\n",
    "        self.ffn = tf.keras.Sequential([Dense(ff_dim, activation=\"relu\"), Dense(embed_dim),])\n",
    "        self.layernorm1 = LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = LayerNormalization(epsilon=1e-6)\n",
    "        self.dropout1 = Dropout(rate)\n",
    "        self.dropout2 = Dropout(rate)\n",
    "    def call(self,inputs,training):\n",
    "        attn_output = self.att(inputs[0],inputs[1])\n",
    "        attn_output = self.dropout1(attn_output,training=training)\n",
    "        out1 = self.layernorm1(inputs[0]+attn_output)\n",
    "        ffn_output = self.ffn(out1)\n",
    "        ffn_output = self.dropout2(ffn_output, training=training)\n",
    "        return self.layernorm2(out1+ffn_output)\n",
    "def get_model(input_shape,output_shape,tokenizer,embed_dim=32,p1=[[8,32],[16,32]],p2=[[16,32],[8,32]]):\n",
    "    x_position_embedding = Embedding(input_dim=input_shape[0],output_dim=embed_dim)(np.array([i for i in range(input_shape[0])]).reshape(1,-1))\n",
    "    x = Input(shape=input_shape)\n",
    "    input_embedding = Embedding(input_dim=len(tokenizer.word_index)+1,output_dim=embed_dim)(x)\n",
    "    knowledge = input_embedding+x_position_embedding\n",
    "    for i in p1:\n",
    "        knowledge = TransformerBlock(embed_dim,i[0],i[1])(knowledge)+x_position_embedding\n",
    "    y_position_embedding = Embedding(input_dim=output_shape[0],output_dim=embed_dim)(np.array([i for i in range(output_shape[0])]).reshape(1,-1))\n",
    "    answer = tf.tile(knowledge[:,:1,:],[1,output_shape[0],1])*0+y_position_embedding\n",
    "    for i in p2:\n",
    "        answer = TransformerBlock2(embed_dim,i[0],i[1])([answer,knowledge])+y_position_embedding\n",
    "    y = Dense(len(tokenizer.word_index)+1,activation=\"softmax\")(answer)\n",
    "    return Model(inputs=x,outputs=y)\n",
    "class CX_LOSS(tf.losses.Loss):\n",
    "    def __init__(self, num_classes=80, alpha=0.25, gamma=2.0, delta=1.0):\n",
    "        super().__init__(reduction=\"auto\", name=\"RetinaNetLoss\")\n",
    "        self.l = tf.keras.losses.CategoricalCrossentropy()\n",
    "    def call(self,y_true,y_pred):\n",
    "        loss = self.l(tf.one_hot(y_true,len(tokenizer.word_index)+1),y_pred)\n",
    "        loss = tf.where(y_true==0,0.1*loss,5*loss)\n",
    "        return tf.reduce_sum(loss,axis=-1)\n",
    "model = get_model((100,),(300,),tokenizer,embed_dim=128,p1=[[4,128],[8,128],[16,128],[32,128]],p2=[[32,128],[16,128],[8,128],[4,128]])\n",
    "model.compile(optimizer=Adam(learning_rate=1e-4),loss=CX_LOSS(),metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('inducer2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('inducer2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_g = data_generator(32)\n",
    "history = model.fit_generator(data_g,steps_per_epoch=100,epochs=500)\n",
    "np.save(\"nrt_loss\",history.history['loss'])\n",
    "np.save(\"nrt_history\",history.history['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equation = generate_expression(3,100)\n",
    "equation = '((56+28)+(81*45)*(47+94)+(53*36))'\n",
    "print(f'question: {equation}=?')\n",
    "print(f'\\n')\n",
    "print(f'computation:')\n",
    "compute(equation,False)\n",
    "print(f'\\n')\n",
    "print(f'ground truth = {eval(equation)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "CX_universe",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
